﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LowLevelGraphics.Filter
{
    /// <summary>
    /// Gradient Domain High Dynamic Range Compression
    /// Raanan Fattal, Dani Lischinski, Michael Werman
    /// 
    /// NOTE: minimum image size is MIN_LEVEL_SIZE
    /// </summary>
    /// 
    public sealed class Fattal : Operator
    {
        private const int MIN_LEVEL_SIZE = 32;
        private const float ALPHA = 0.1f;
        private const float EPSILON = 0.0001f;
        private const int MULTIGRID_CYCLES = 3;

        public double Alpha
        {

            get
            {
                return this.alpha;
            }
            set
            {
                this.alpha = Math.Min(Math.Max(value, 0.1), 0.9);
            }
        }

        public double Beta
        {

            get
            {
                return this.beta;
            }
            set
            {
                this.beta = Math.Min(Math.Max(value, 0.6), 0.9);
            }
        }

        public double Saturation
        {

            get
            {
                return this.saturation;
            }
            set
            {
                this.saturation = Math.Min(Math.Max(value, 0.2), 0.6);
            }
        }

        private double alpha = 0.1;
        private double beta = 0.8;
        private double saturation = 0.4;

        public Fattal(double beta, double saturation)
        {

            this.Beta = beta;
            this.Saturation = saturation;
        }

        public Fattal()
        {
        }

        protected override void Run(float[][,] mapInput, int width, int height, ref float[][,] mapOutput)
        {

            // compute normalized luminance map
            float[,] mapLumIn = ComputeLuminanceMap(mapInput, width, height);
            float[,] mapLum = ComputeLogLuminanceMap(mapLumIn, width, height);

            if (mapLum == null)
                return;

            // compute gaussian pyramid
            float[][,] pyramid = ComputeGaussianPyramid(mapLum, width, height);

            // compute gradient magnitude
            float[][,] gradients = null;
            float[] avgGrads = null;

            ComputeGradientPyramid(pyramid, ref gradients, ref avgGrads);

            // free the gaussian pyramid
            pyramid = null;

            GC.Collect();

            // compute the Gradient attenuation function
            float[,] phi = ComputePhi(gradients, avgGrads);

            // free the gradient pyramid
            gradients = null;
            avgGrads = null;

            GC.Collect();

            // compute divergence
            float[,] div = ComputeDivergence(mapLum, phi);

            // free mapLum and phi
            // TODO: uncomment
            mapLum = null;
            phi = null;

            GC.Collect();

            // solve PDE
            MultigridSolver solver = new MultigridSolver();

            int size = solver.getSize(width, height);
            int x = 0, y = 0;

            double[,] Id = new double[size, size];

            // shift pixels to create a boundary
            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    Id[x + 1, y + 1] = ((x < width && y < height) ? div[x, y] : 0);

            //SaveArray(I, "multigrid_before");

            solver.mglin(Id, size, MULTIGRID_CYCLES);

            // shift pixels back
            double min = Double.MaxValue;
            double max = Double.MinValue;

            for (y = 1; y <= height; y++)
                for (x = 0; x <= width; x++)
                {

                    if (Id[x, y] < min)
                        min = Id[x, y];

                    if (Id[x, y] > max)
                        max = Id[x, y];
                }

            float[,] I = new float[width, height];

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    I[x, y] = (float)((Id[x + 1, y + 1] - min) / (max - min));

            // exp luminance map
            ExpLuminanceMap(I, width, height);

            // normalize luminance map
            NormalizeLuminanceMap(I, width, height);

            // compute output
            ComputeOutputMap(mapInput, mapLumIn, I, width, height, ref mapOutput);
        }

        protected override void SetParam(string param, string value)
        {

            switch (param)
            {
                case "alpha":
                    this.Alpha = Double.Parse(value);
                    break;
                case "beta":
                    this.Beta = Double.Parse(value);
                    break;
                case "saturation":
                    this.Saturation = Double.Parse(value);
                    break;
            }
        }

        private float[,] ComputeLuminanceMap(float[][,] mapInput, int width, int height)
        {

            int x = 0, y = 0;
            float[,] mapLum = new float[width, height];

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    mapLum[x, y] = (float)(0.2992 * mapInput[2][x, y] + 0.5868 * mapInput[1][x, y] + 0.1140 * mapInput[0][x, y]);

            return mapLum;
        }

        private float[,] ComputeLogLuminanceMap(float[,] mapLum, int width, int height)
        {

            int x = 0, y = 0;
            float lum = 0;
            float max = Single.MinValue; // maximum luminance
            float min = Single.MaxValue; // minimum luminance
            float[,] mapLogLum = new float[width, height];

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                {

                    lum = mapLum[x, y];
                    max = ((lum > max) ? lum : max);
                    min = ((lum < min) ? lum : min);
                }

            if (max == min)
                return null;

            // normalize to <0,100> and compute logarithm
            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    mapLogLum[x, y] = (float)Math.Log(EPSILON + 100 * (mapLum[x, y] - min) / (max - min));

            return mapLogLum;
        }

        private float[][,] ComputeGaussianPyramid(float[,] mapLum, int width, int height)
        {

            // get the number of pyramid levels
            int minSize = Math.Min(width, height);
            int levels = 0;

            while (minSize >= MIN_LEVEL_SIZE)
            {

                levels++;
                minSize >>= 1;
            }

            // create the Gaussian pyramid
            float[][,] pyramid = new float[levels][,];

            // clone first level
            int x = 0, y = 0;

            pyramid[0] = new float[width, height];

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    pyramid[0][x, y] = mapLum[x, y];

            // create higher levels
            for (int i = 1; i < levels; i++)
                pyramid[i] = GaussianLevel5x5(pyramid[i - 1]);

            return pyramid;
        }

        private void ComputeGradientPyramid(float[][,] pyramid, ref float[][,] gradients, ref float[] avgGrads)
        {

            int levels = pyramid.Length;

            gradients = new float[levels][,];
            avgGrads = new float[levels];

            for (int i = 0; i < levels; i++)
            {

                float[,] gradLevel = null;
                float avgGrad = 0;

                GradientLevel(pyramid[i], i, ref gradLevel, ref avgGrad);

                gradients[i] = gradLevel;
                avgGrads[i] = avgGrad;
            }
        }

        private float[,] ComputePhi(float[][,] gradients, float[] avgGrads)
        {

            int levels = gradients.Length;

            float[][,] phi = new float[levels][,];

            float[,] gradLevel = null;
            int width = 0, height = 0;
            float alpha = 0;
            int x = 0, y = 0;

            for (int i = (levels - 1); i >= 0; i--)
            {

                gradLevel = gradients[i];
                width = gradLevel.GetLength(0);
                height = gradLevel.GetLength(1);

                alpha = (ALPHA * avgGrads[i] * (1 << i));

                if (alpha == 0)
                    alpha = EPSILON;

                phi[i] = new float[width, height];

                for (y = 0; y < height; y++)
                    for (x = 0; x < width; x++)
                        phi[i][x, y] = (float)Math.Min(Math.Pow(gradLevel[x, y] / alpha, this.beta - 1), 1);

                if (i < (levels - 1))
                {

                    float[,] l = new float[width, height];

                    // TODO: bilinear filtering
                    for (y = 0; y < height; y++)
                        for (x = 0; x < width; x++)
                            l[x, y] = phi[i + 1][x >> 1, y >> 1];

                    for (y = 0; y < height; y++)
                        for (x = 0; x < width; x++)
                            phi[i][x, y] *= l[x, y];

                    l = null;
                    phi[i + 1] = null;
                }
            }

            return phi[0];
        }

        private float[,] ComputeDivergence(float[,] mapLum, float[,] phi)
        {

            int width = mapLum.GetLength(0);
            int height = mapLum.GetLength(1);

            // gradient attenuation
            float[,] gx = new float[width, height];
            float[,] gy = new float[width, height];

            int x = 0, y = 0;
            int s = 0, e = 0;

            for (y = 0; y < height; y++)
            {

                s = ((y == (height - 1)) ? y : (y + 1));

                for (x = 0; x < width; x++)
                {

                    e = ((x == (width - 1)) ? x : (x + 1));

                    gx[x, y] = ((mapLum[e, y] - mapLum[x, y]) * phi[x, y]);
                    gy[x, y] = ((mapLum[x, s] - mapLum[x, y]) * phi[x, y]);
                }
            }

            // divergence
            float[,] div = new float[width, height];

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                {

                    div[x, y] = (gx[x, y] + gy[x, y]);

                    if (x > 0)
                        div[x, y] -= gx[x - 1, y];

                    if (y > 0)
                        div[x, y] -= gy[x, y - 1];
                }

            gx = null;
            gy = null;

            return div;
        }

        private void ExpLuminanceMap(float[,] mapLum, int width, int height)
        {

            int x = 0, y = 0;

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    mapLum[x, y] = (float)(Math.Exp(mapLum[x, y]) - EPSILON);
        }

        private void NormalizeLuminanceMap(float[,] mapLum, int width, int height)
        {

            int x = 0, y = 0;
            float lum = 0;
            float min = Single.MaxValue;
            float max = Single.MinValue;

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                {

                    lum = mapLum[x, y];

                    max = ((lum > max) ? lum : max);
                    min = ((lum < min) ? lum : min);
                }

            for (y = 0; y < height; y++)
                for (x = 0; x < width; x++)
                    mapLum[x, y] = Math.Min(Math.Max(((mapLum[x, y] - min) / (max - min)), 0), 1);
        }

        private void ComputeOutputMap(float[][,] mapInput, float[,] mapLumIn, float[,] mapLum, int width, int height, ref float[][,] mapOutput)
        {

            int x = 0, y = 0;

            for (int c = 0; c < 3; c++)
                for (y = 0; y < height; y++)
                    for (x = 0; x < width; x++)
                        mapOutput[c][x, y] = (float)Math.Min(Math.Max((Math.Pow(mapInput[c][x, y] / mapLumIn[x, y], this.saturation) * mapLum[x, y]), 0), 1);
        }

        private float[,] GaussianLevel5x5(float[,] level)
        {

            int width = level.GetLength(0);
            int height = level.GetLength(1);
            int nWidth = ((width >> 1) + (width % 2));
            int nHeight = ((height >> 1) + (height % 2));

            int x = 0, y = 0;
            float[,] tempLevelH = new float[width, height];
            float[,] tempLevelV = new float[width, height];
            float[,] newLevel = new float[nWidth, nHeight];

            // horizontal filtering
            for (y = 0; y < height; y++)
            {

                for (x = 2; x < (width - 2); x++)
                    tempLevelH[x, y] = ((level[x - 2, y] + level[x + 2, y] + 4 * (level[x - 1, y] + level[x + 1, y]) + 6 * level[x, y]) / 16);

                tempLevelH[0, y] = ((2 * level[2, y] + 8 * level[1, y] + 6 * level[0, y]) / 16);
                tempLevelH[1, y] = ((level[3, y] + 4 * (level[0, y] + level[2, y]) + 7 * level[1, y]) / 16);
                tempLevelH[width - 2, y] = ((level[width - 4, y] + 5 * level[width - 1, y] + 4 * level[width - 3, y] + 6 * level[width - 2, y]) / 16);
                tempLevelH[width - 1, y] = ((level[width - 3, y] + 5 * level[width - 2, y] + 10 * level[width - 1, y]) / 16);
            }

            // vertical filtering
            for (x = 0; x < width; x++)
            {

                for (y = 2; y < (height - 2); y++)
                    tempLevelV[x, y] = ((tempLevelH[x, y - 2] + tempLevelH[x, y + 2] + 4 * (tempLevelH[x, y - 1] + tempLevelH[x, y + 1]) + 6 * tempLevelH[x, y]) / 16);

                tempLevelV[x, 0] = ((2 * tempLevelH[x, 2] + 8 * tempLevelH[x, 1] + 6 * tempLevelH[x, 0]) / 16);
                tempLevelV[x, 1] = ((tempLevelH[x, 3] + 4 * (tempLevelH[x, 0] + tempLevelH[x, 2]) + 7 * tempLevelH[x, 1]) / 16);
                tempLevelV[x, height - 2] = ((tempLevelH[x, height - 4] + 5 * tempLevelH[x, height - 1] + 4 * tempLevelH[x, height - 3] + 6 * tempLevelH[x, height - 2]) / 16);
                tempLevelV[x, height - 1] = ((tempLevelH[x, height - 3] + 5 * tempLevelH[x, height - 2] + 10 * tempLevelH[x, height - 1]) / 16);
            }

            // resampling
            // TODO: bilinear filtering
            for (y = 0; y < nHeight; y++)
                for (x = 0; x < nWidth; x++)
                    newLevel[x, y] = tempLevelV[x << 1, y << 1];

            tempLevelH = null;
            tempLevelV = null;

            return newLevel;
        }

        private void GradientLevel(float[,] level, int k, ref float[,] gradLevel, ref float avgGrad)
        {

            int width = level.GetLength(0);
            int height = level.GetLength(1);
            float divider = (float)(1 << (k + 1));

            int x = 0, y = 0, n = 0, s = 0, w = 0, e = 0;
            float gx = 0, gy = 0, mag = 0;
            gradLevel = new float[width, height];
            avgGrad = 0;

            for (y = 0; y < height; y++)
            {

                n = ((y == 0) ? 0 : (y - 1));
                s = ((y == (height - 1)) ? y : (y + 1));

                for (x = 0; x < width; x++)
                {

                    w = ((x == 0) ? 0 : (x - 1));
                    e = ((x == (width - 1)) ? 0 : (x + 1));

                    gx = ((level[e, y] - level[w, y]) / divider);
                    gy = ((level[x, s] - level[x, n]) / divider);

                    mag = (float)Math.Sqrt(gx * gx + gy * gy);

                    gradLevel[x, y] = mag;
                    avgGrad += mag;
                }
            }

            avgGrad /= (width * height);
        }
    }
}
